import os, torch
import numpy as np
from numba import njit, jit
from UTILS.colorful import *
from .gcortex import GNet
from config import GlobalConfig
from UTILS.tensor_ops import copy_clone, my_view, add_onehot_id_at_last_dim, add_obs_container_subject, __hash__
DEBUG = True

class AlgorithmConfig:
    n_cluster = 6
    num_cluster = 6
    gamma = 0.99
    tau = 0.95
    train_traj_needed = 512
    upper_training_epoch = 4
    load_checkpoint = False
    take_reward_as_unity = True

    class Shell_Config:
        max_internal_step = 5
        head_start_cnt = 3
        decision_interval = 30
        head_start_hold_n = 15

    class PPO_Config:
        clip_param = 0.2
        ppo_epoch = 16
        n_pieces_batch_division = 3
        value_loss_coef = 0.1
        entropy_coef = 0.05
        max_grad_norm = 0.5
        clip_param = 0.2
        lr = 1e-4
        balance = 0.5
        add_prob_loss = True

class ReinforceAlgorithmFoundation(object):
    def __init__(self, n_agent, n_thread, space, mcv=None):
        self.n_thread = n_thread
        self.n_agent = n_agent
        self.act_space = space['act_space']
        self.obs_space = space['obs_space']
        scenario_config = GlobalConfig.scenario_config
        alg_config = AlgorithmConfig
        from .shell_env import ShellEnvWrapper
        self.shell_env = ShellEnvWrapper(n_agent, n_thread, space, mcv, self, 
                                        alg_config, scenario_config)
        self.device = GlobalConfig.device
        self.policy = GNet(num_agents=self.n_agent, 
                           num_entities=scenario_config.num_entity, 
                           num_cluster=alg_config.num_cluster,
                           n_basic_dim=scenario_config.obs_vec_length).to(self.device)
        self.AvgRewardAgentWise = alg_config.take_reward_as_unity
        from .ppo import PPO
        self.trainer = PPO(self.policy, ppo_config=AlgorithmConfig.PPO_Config, mcv=None)
        from .trajectory import BatchTrajManager
        self.batch_traj_manager = BatchTrajManager(n_env=n_thread, 
                                                   traj_limit=int(scenario_config.max_steps_episode*1.5), 
                                                   templete='off-policy',
                                                   trainer_hook=self.trainer.train_on_traj,
                                                   purge_outdated=True)

        self.load_checkpoint = AlgorithmConfig.load_checkpoint
        note = GlobalConfig.note
        if not os.path.exists('./ZHECKPOINT/%s/'%note): 
            os.makedirs('./ZHECKPOINT/%s/'%note)
        if self.load_checkpoint:
            print黄('加载检查点')
            self.policy.load_state_dict(torch.load('%s/model.pt'%logdir, map_location=GlobalConfig.cuda_device))
        self.__incomplete_rollout_frag__ = None
        self.patience = 500 # skip currupt data detection after patience exhausted

    # _____________________Redirection____________________
    # this is a redirect to shell_env.interact_with_env
    def interact_with_env(self, State_Recall):
        return self.shell_env.interact_with_env(State_Recall)

    # self.interact_with_env --> shell_env.interact_with_env --> self.interact_with_env_genuine
    def interact_with_env_genuine(self, State_Recall):
        test_mode = State_Recall['Test-Flag']
        if not test_mode: self.train()
        return self.action_making(State_Recall, test_mode) # state_recall dictionary will preserve states for next action making


    def action_making(self, State_Recall, test_mode):
        # In case of multi-process align call, do nothing and just wait reward
        if State_Recall['all_emb'] is None: return self._interact_empty(State_Recall)
        # this is the real Reinforce processing
        threads_active_flag = State_Recall['threads_active_flag']
        all_emb = State_Recall['all_emb']
        # mask out obs on active threads
        Active_emb = {}
        for k in all_emb: Active_emb[k] = all_emb[k][threads_active_flag]
        # make decision
        with torch.no_grad():
            # act, value, actLogProbs
            Active_action, Active_value, Active_action_log_prob = self.policy.act(Active_emb, test_mode=test_mode)
        # commit action to memory
        wait_reward_hook = self.commit_frag_hook({
            '_SKIP_':            ~threads_active_flag, # thread mask
            'value':            Active_value,
            'g_actionLogProbs': Active_action_log_prob,
            'g_obs':              Active_emb, 
            'g_actions':          Active_action,
        }, require_hook = True) if not test_mode else self.__dummy_hook
        '''   <1>  we will deal with rollout later after the reward is ready, 
                        now we leave a hook to be callback    '''
        State_Recall['_hook_'] = wait_reward_hook
        return Active_action.copy(), State_Recall

    def commit_frag_hook(self, f1, require_hook = True):
        assert self.__incomplete_rollout_frag__ is None
        self.__incomplete_rollout_frag__ = f1 
        self.__check_data_hash() # this is important!
        if require_hook: # leave hook
            return lambda f2: self.rollout_frag_hook(f2)
        return


    # ________Rollout Processor_______
    def rollout_frag_hook(self, f2):
        '''   <2>  hook is called, reward and next moment observation is ready,
                        now feed them into trajectory manager    '''
        # do data curruption check at beginning, this is important!
        self.__check_data_curruption()
        # strip info, since it is not array
        items_to_pop = ['info', 'Latest-Obs']
        strip_dict = {}
        for k in items_to_pop:
            if k in f2: 
                strip_dict[k] = f2.pop(k)
        # the agent-wise reward is supposed to be the same, so averge them
        if self.AvgRewardAgentWise:
            f2['reward'] = np.mean(f2['reward'], axis=-1, keepdims=True)
        # change the name of done to be recognised
        f2['_DONE_'] = f2.pop('done')
        f2['_TOBS_'] = f2.pop('Terminal-Obs-Echo') if 'Terminal-Obs-Echo' in f2 else None
        # integrate frag part1 and part2
        self.__incomplete_rollout_frag__.update(f2)
        completed_frag = self.mask_out_paused(self.__incomplete_rollout_frag__)
        # put the fragment into memory
        self.batch_traj_manager.feed_traj(completed_frag)
        self.__incomplete_rollout_frag__ = None


    def mask_out_paused(self, fragment):
        running = ~fragment['_SKIP_']
        if running.all():
            return fragment
        for key in fragment:
            if not key.startswith('_') and hasattr(fragment[key], '__len__') and len(fragment[key]) == self.n_thread:
                fragment[key] = fragment[key][running]
        return fragment

    def train(self):
        if self.batch_traj_manager.can_exec_training():  # time to start a training routine
            # print('self.decision_interval', self.decision_interval)
            update_cnt = self.batch_traj_manager.train_and_clear_traj_pool()
            self.__save_model(update_cnt)





    def _interact_empty(self, State_Recall):
        test_mod = State_Recall['Test-Flag']
        threads_active_flag = State_Recall['threads_active_flag']
        wait_reward_hook = self.commit_frag_hook({
            '_SKIP_': ~threads_active_flag,  'g_obs': {'agent_final_emb':None,'entity_final_emb':None,'cluster_final_emb':None}, 
            'value': None, 'g_actionLogProbs': None, 'g_actions': None
            }, require_hook = True) if not test_mod else self.__dummy_hook
        return None, {'_hook_':wait_reward_hook}
























    def __dummy_hook(self, f2): 
        return

    def __save_model(self, update_cnt):
        note = GlobalConfig.note
        if os.path.exists('./ZHECKPOINT/%s/save_now.txt'%note) or update_cnt%100==99:
            print绿('保存模型中')
            torch.save(self.policy.state_dict(), '%s/model.pt'%logdir)
            print绿('保存模型完成')
            try:
                os.remove('./ZHECKPOINT/%s/save_now.txt'%note)
            except:
                print('清除保存标志:异常')

    # debugging functions
    def __check_data_hash(self):
        if self.patience > 0: 
            self.hash_debug = {}
            # for debugging, to detect write protection error
            for key in self.__incomplete_rollout_frag__:
                item = self.__incomplete_rollout_frag__[key]
                if isinstance(item, dict):
                    self.hash_debug[key]={}
                    for subkey in item:
                        subitem = item[subkey]
                        self.hash_debug[key][subkey] = __hash__(subitem)
                else:
                    self.hash_debug[key] = __hash__(item)

    def __check_data_curruption(self):
        if self.patience > 0: 
            assert self.__incomplete_rollout_frag__ is not None
            assert self.hash_debug is not None
            for key in self.__incomplete_rollout_frag__:
                item = self.__incomplete_rollout_frag__[key]
                if isinstance(item, dict):
                    for subkey in item:
                        subitem = item[subkey]
                        assert self.hash_debug[key][subkey] == __hash__(subitem), ('Currupted data! 发现腐败数据!')
                else:
                    assert self.hash_debug[key] == __hash__(item), ('Currupted data! 发现腐败数据!')
            self.patience -= 1
